import ppq.lib as PFL
import torch
import torchvision
from ppq.api import ENABLE_CUDA_KERNEL, load_native_graph, load_torch_model
from ppq.core import TargetPlatform
from ppq.executor import TorchExecutor
from ppq.quantization.optim import *
from ppq.quantization.quantizer import TensorRTQuantizer
from ppq.samples.Imagenet.Utilities.Imagenet import *  # check ppq.samples.imagenet.Utilities
from ppq.samples.Imagenet.Utilities.Imagenet.imagenet_util import \
    load_imagenet_from_directory  # check ppq.samples.imagenet.Utilities

from trainer_train import ImageNetTrainer

"""
    使用这个脚本来尝试在 Imagenet 数据集上执行量化感知训练
        使用 imagenet 中的数据测试量化精度与 calibration
        默认的 imagenet 数据集位置:Assets/Imagenet_Train, Assets/Imagenet_Valid
        你可以通过软连接创建它们:
            ln -s /home/data/Imagenet/val Assets/Imagenet_Valid
            ln -s /home/data/Imagenet/train Assets/Imagenet_Train
"""

CFG_DEVICE = 'cuda'  # 一个神奇的字符串，用来确定执行设备
CFG_BATCHSIZE = 64  # 测试与calib时的 batchsize
CFG_INPUT_SHAPE = (CFG_BATCHSIZE, 3, 224, 224)  # 用来确定模型输入的尺寸，好像 imagenet 都是这个尺寸
CFG_VALIDATION_DIR = '../../../datasets/tiny-imagenet-200/val'  # 用来读取 validation dataset
CFG_TRAIN_DIR = '../../../datasets/tiny-imagenet-200/train'  # 用来读取 train dataset，注意该集合将被用来 calibrate 你的模型
CFG_PLATFORM = TargetPlatform.TRT_INT8  # 用来指定目标平台

# ------------------------------------------------------------
# 在这个例子中我们将向你展示如何在 PPQ 中对你的网络进行量化感知训练
# 你可以使用带标签的数据执行正常的训练流程，也可以使用类似蒸馏的方式进行无标签训练
# PPQ 模型的训练过程与 Pytorch 遵循相同的逻辑，你可以使用 Pytorch 中的技巧来获得更好的训练效果
# ------------------------------------------------------------
model = torchvision.models.mobilenet.mobilenet_v2(weights=torchvision.models.mobilenet.MobileNet_V2_Weights.IMAGENET1K_V2).to("cuda")
# model = model.load_state_dict(torch.load("Best.native")).to("cuda")
# weights=torch.load("Best.native")
calib_dataloader = load_imagenet_from_directory(
    directory=CFG_TRAIN_DIR, batchsize=CFG_BATCHSIZE,
    shuffle=True, subset=1280, require_label=True,
    num_of_workers=0)

training_dataloader = load_imagenet_from_directory(
    directory=CFG_TRAIN_DIR, batchsize=CFG_BATCHSIZE,
    shuffle=True, require_label=True,
    num_of_workers=0)

eval_dataloader = load_imagenet_from_directory(
    directory=CFG_VALIDATION_DIR, batchsize=CFG_BATCHSIZE,
    shuffle=True, require_label=True,
    num_of_workers=0)

# with ENABLE_CUDA_KERNEL():
# ------------------------------------------------------------
# 创建优化管线，由于后续还要继续训练我们的模型，我们不能在此处调用
# ParameterBakingPass()，一旦模型权重完成烘焙，则它们不能被进一步调整
# ------------------------------------------------------------

trainer = ImageNetTrainer(model=model)

best_acc = 0
trainer.eval(eval_dataloader)
for epoch in range(20):
    # print(1)
    trainer.epoch(training_dataloader)
    current_acc = trainer.eval(eval_dataloader)
    if current_acc > best_acc:
        trainer.save('Best.native')
